import cv2
import matplotlib.pyplot as plt
import numpy as np

#  退化图像的维纳滤波 (Wiener filter)
def getMotionDsf(img, angle, dist):
    xCenter = img.shape[0]// 2
    yCenter = img.shape[1]// 2
    sinVal = np.sin(angle * np.pi / 180)
    cosVal = np.cos(angle * np.pi / 180)
# 点扩散函数
    PSF = np.zeros(img.shape[:2])
# 将对应角度上motion_dis个点置成1
    for i in range(dist):
        xOffset = round(sinVal * i)
        yOffset = round(cosVal * i)
        PSF[int(xCenter - xOffset), int(yCenter + yOffset)] = 1
 # 归一化
    return PSF / PSF.sum()


# 维纳滤波，K默认为0.01
def wienerFilter(img, PSF, eps, K=0.005):
    fftImg = np.fft.fft2(img)
    fftPSF = np.fft.fft2(PSF) + eps
# np.conj返回共轭
    fftWiener = np.conj(fftPSF) / (np.abs(fftPSF)**2 + K)
    imgWienerFilter = np.fft.ifft2(fftImg * fftWiener)
    imgWienerFilter = np.abs(np.fft.fftshift(imgWienerFilter))
    return imgWienerFilter

# 返回 P(u,v)
def getPuv(img):
    h, w = img.shape[:2]
    hPad, wPad = h - 3, w - 3
    pxy = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
    pxyPad = np.pad(pxy, ((hPad // 2, hPad - hPad // 2), (wPad // 2, wPad - wPad // 2)), mode='constant')
    fftPuv = np.fft.fft2(pxyPad)
    return fftPuv

# 约束最小二乘方滤波
def leastSquareFilter(img, PSF, eps, gamma=0.01):
    fftImg = np.fft.fft2(img)
    fftPSF = np.fft.fft2(PSF)
    conj = fftPSF.conj()
    fftPuv = getPuv(img)

    Huv = conj / (np.abs(fftPSF)**2 + gamma * (np.abs(fftPuv)**2))
    ifftImg = np.fft.ifft2(fftImg * Huv)
    ifftShift = np.abs(np.fft.fftshift(ifftImg))
    imgLSFilter = np.uint8(cv2.normalize(np.abs(ifftShift), None, 0, 255, cv2.NORM_MINMAX)) # 归一化为 [0,255]
    return imgLSFilter


# 读取原始图像
img = cv2.imread("D:\\pycharmproject\\pythonProject\\order\\boy-blurred(1).tif", 0) # flags=0 读取为灰度图像

# 估计点扩散函数 (PSF)
# 修改角度和距离参数以匹配实际模糊情况
PSF = getMotionDsf(img, angle=50, dist=89)

# 去除模糊：逆滤波与维纳滤波
imgWienerFilter = wienerFilter(img, PSF, eps=1e-6, K=0.0008)
imgLSFilter = leastSquareFilter(img, PSF, eps=1e-6, gamma=0.05) # 0.1

def enhance_deblurred_image(img):
    # 对去模糊后的图像进行增强
    # 对比度增强
    clahe = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(1, 1))
    enhanced = clahe.apply(img)

    # 锐化
    kernel = np.array([[-1, -1, -1],
                       [-1, 9.4, -1],
                       [-1, -1, -1]])
    sharpened = cv2.filter2D(enhanced, -1, kernel)

    return sharpened

# def enhance_deblurred_image(img):
#     # 自适应直方图均衡
#     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(1, 1))
#     enhanced = clahe.apply(img)
#
#     # 边缘增强锐化
#     kernel = np.array([
#         [0, -1, 0],
#         [-1, 5.1, -1],
#         [0, -1, 0]
#     ])
#     sharpened = cv2.filter2D(enhanced, -1, kernel)
#
#     # 对比度拉伸
#     min_val, max_val = np.min(sharpened), np.max(sharpened)
#     stretched = cv2.normalize(sharpened, None, 0, 255, cv2.NORM_MINMAX)
#
#     return np.uint8(stretched)
#
# def enhance_deblurred_image(img):
#     # 双边滤波去噪，保留边缘
#     denoised = cv2.bilateralFilter(img, 9, 75, 75)
#
#     # 自适应直方图均衡
#     clahe = cv2.createCLAHE(clipLimit=2, tileGridSize=(3, 3))
#     enhanced = clahe.apply(denoised)
#
#     # 锐化
#     kernel = np.array([
#         [0, -1, 0],
#         [-1, 5.4, -1],
#         [0, -1, 0]
#     ])
#     sharpened = cv2.filter2D(enhanced, -1, kernel)
#
#     return sharpened

# def enhance_deblurred_image(img):
#     # 多步骤增强
#
#     # 第一步：去噪
#     denoised = cv2.fastNlMeansDenoising(img, None, 9, 9, 11)
#
#     # 第二步：自适应直方图均衡
#     clahe = cv2.createCLAHE(clipLimit=0.5, tileGridSize=(2, 2))
#     enhanced = clahe.apply(denoised)
#
#     # 第三步：复杂锐化
#     kernel = np.array([
#         [-1 / 9, -1 / 9, -1 / 9],
#         [-1 / 9, 2.1, -1 / 9],
#         [-1 / 9, -1 / 9, -1 / 9]
#     ])
#     sharpened = cv2.filter2D(enhanced, -1, kernel)
#
#     # 第四步：对比度拉伸
#     stretched = cv2.normalize(sharpened, None, 0, 255, cv2.NORM_MINMAX)
#
#     # 第五步：伽马校正
#     gamma_corrected = np.array(255 * (stretched / 255) ** 0.7, dtype=np.uint8)
#
#     return gamma_corrected

imgLSFilter_enhanced = enhance_deblurred_image(imgLSFilter)
imgWienerFilter_enhanced = enhance_deblurred_image(imgWienerFilter.astype(np.uint8))

# 显示结果
plt.figure(figsize=(12, 8))
plt.subplot(1, 3, 1)
plt.imshow(img, cmap='gray')  # 使用 'gray' 而不是 'viridis'
plt.title('Original Image')
plt.axis('off')

plt.subplot(1, 3, 2)
plt.imshow(imgWienerFilter_enhanced, cmap='gray')
plt.title('Wiener Filter')
plt.axis('off')

plt.subplot(1, 3, 3)
plt.imshow(imgLSFilter_enhanced, cmap='gray')
plt.title('Least Square Filter')
plt.axis('off')

plt.tight_layout()
plt.show()  # 添加这一行来显示图像